import jax.numpy as jnp
from flax import nnx

from model import EnsembleFKModel, MlpPartialIKModel
from model.data_processing.data_loaders import Batch, NormalizedEndJointData
from model.data_processing.jax_dataloader import NumpyLoader

"""
Training a forward kinematics model
Joints -> End pos
"""


class NormalizedFKData(NormalizedEndJointData):
    """
    Given 6D joints -> predict 3D end positions
    """

    def __init__(self, filename, arm_end_length: float = 130):
        super().__init__(filename, arm_end_length)

    def __getitem__(self, idx):
        return Batch(self.joints[idx], self.ends[idx][:3])


class FKvalue(nnx.Module):
    """
    evaluate the distante(fk(joints), target_xyz)
    """

    def __init__(self, fk_model: nnx.Module):
        self.fk_model = fk_model

    def __call__(self, joints, target_xyz):
        mse = (self.fk_model(joints) - target_xyz) ** 2

        return jnp.sum(mse, axis=-1)


if __name__ == "__main__":
    data = NormalizedFKData("data/json_data/all.json")
    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    # reuse the similar structure as the end-joint model
    # model = EnsembleFKModel(in_dim=6, out_dim=3, num_heads=10, weight_decay=1e-4)
    model = MlpPartialIKModel(in_dim=6, out_dim=3, weight_decay=1e-2)

    model.train(data_loader, 5000, save_as="mlp_normed_FK.ckp")
